import numpy as np
import cv2
import cv2.aruco as aruco
import os

def special_resize(A_i, B_i):
    a = np.zeros(shape=(B_i.shape[0], 9))
    for i in range(B_i.shape[0]):
        for j in range(B_i.shape[0]):
            a[i, j+i*A_i.shape[0]] = A_i[j]
    b = B_i.reshape((B_i.shape[0], -1))
    return a, b

class CalibrateNinePointsTakeAPhoto(object):
    """
    针尖校准（拍照1次）
    Robot: 机器人用工具针尖对准9个Mark点（工具坐标）
    Camera: 相机在固定点位同时拍9个Mark点（图像坐标）
    """
    def __init__(self, board_size):
        # 找棋盘格角点
        self.criteria_findCorner = (cv2.CALIB_CB_ADAPTIVE_THRESH + cv2.CALIB_CB_FAST_CHECK + cv2.CALIB_CB_NORMALIZE_IMAGE)
        # 设置寻找亚像素角点的参数，采用的停止准则是最大循环次数30和最大误差容限0.001
        self.criteria_cornerSubPix = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001) # 阈值
        # 棋盘格模板规格
        self.board_size = board_size

    def get_nine_points(self, img_path, point_indexs):
        img = cv2.imread(img_path)
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        # 找到棋盘格角点
        ret, corners = cv2.findChessboardCorners(gray, (self.board_size[1], self.board_size[0]), self.criteria_findCorner)
        # 如果找到足够点对，将其存储起来
        if ret == True:
            # 在原角点的基础上寻找亚像素角点
            cv2.cornerSubPix(gray, corners, (11,11), (-1,-1), self.criteria_cornerSubPix)
            points = corners[point_indexs, 0, :]
            self.show_points(img, points)

            # 将角点在图像上显示
            # cv2.drawChessboardCorners(img, (self.board_size[1], self.board_size[0]), corners, ret)
            # cv2.namedWindow('findCorners', cv2.WINDOW_NORMAL)
            # cv2.resizeWindow('findCorners', 640, 480)
            # cv2.imshow('findCorners', img)
            # cv2.waitKey(0)
        return points

    def show_points(self, img, points):
        for p in points:
            cv2.circle(img, (int(p[0]),int(p[1])), 7, (0,0,255), 3)
        cv2.imwrite('nine_points.bmp', img)

    def solve_inv(self, A, B):
        A_t = A.T
        X = B @ A_t @ np.linalg.inv(A @ A_t)
        return X
    
    def solve_lstsq(self, A, B):
        """
        应用最小二乘法，需重新构造系数矩阵
        """
        assert A.shape[1]==B.shape[1], "矩阵A与矩阵B的第二维度必须相等!"
        A_list = []
        b_list = []
        for i in range(A.shape[1]):
            Ai, bi = special_resize(A[:,i], B[:,i])
            A_list.append(Ai)
            b_list.append(bi)
        A_ = np.concatenate(A_list, axis=0)
        b_ = np.concatenate(b_list, axis=0)

        x = np.linalg.lstsq(A_, b_, rcond=None)[0]
        return x
    
class CalibrateNinePointsTakeNinePhotos(object):
    """
    相机校准（拍照9次）
    Robot: 机器人带动相机移动9个不同点位（法兰坐标）
    Camera: 相机在不同点位拍同1个Mark点（图像坐标）
    """
    def __init__(self):
        pass

    def get_mark_pixel(self, image):
        # image = cv2.imread(img_path)
        image_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

        # 取ArUco码的四个角点作为特征
        aruco_dict = aruco.getPredefinedDictionary(aruco.DICT_7X7_1000)
        parameters = aruco.DetectorParameters()
        detector = aruco.ArucoDetector(aruco_dict, parameters)
        corners, ids, rejected = detector.detectMarkers(image_gray)

        # print(corners[0][0,0,:])
        # cv2.circle(image, (int(corners[0][0,0,0]), int(corners[0][0,0,1])), 2, (255,0,0), 5)
        # cv2.namedWindow("img",0)
        # cv2.resizeWindow("img", 640, 480)
        # cv2.imshow("img", image)
        # cv2.waitKey(0)

        return corners[0][0,0,:].tolist()
    
    def load_mark_pixels(self, data_path):
        pixels = []
        for i in range(9):
            img_path = os.path.join(data_path, str(i)+".bmp")
            image = cv2.imread(img_path)
            mark_pixel = self.get_mark_pixel(image)
            # mark_pixel.append(1)
            pixels.append(mark_pixel)

        pixels = np.asarray(pixels)

        return pixels
    
    def solve_inv(self, A, B):
        A_t = A.T
        X = B @ A_t @ np.linalg.inv(A @ A_t)
        return X
    
    def solve_lstsq(self, A, B):
        """
        应用最小二乘法，需重新构造系数矩阵
        """
        assert A.shape[1]==B.shape[1], "矩阵A与矩阵B的第二维度必须相等!"
        A_list = []
        b_list = []
        for i in range(A.shape[1]):
            Ai, bi = special_resize(A[:,i], B[:,i])
            A_list.append(Ai)
            b_list.append(bi)
        A_ = np.concatenate(A_list, axis=0)
        b_ = np.concatenate(b_list, axis=0)

        x = np.linalg.lstsq(A_, b_, rcond=None)[0]
        return x


def test_CalibrateNinePointsTakeAPhoto():
    cali = CalibrateNinePointsTakeAPhoto((11,8))

    # 获取像素坐标系下的坐标
    img_path = "D:/VSCodeProjects/calibration/data/img.bmp"
    point_indexs = [1, 4, 7, 41, 44, 47, 81, 84, 87]
    pixels = cali.get_nine_points(img_path, point_indexs)
    
    # 获取法兰坐标系下的坐标
    base_pose = np.array([[-318.74, -858.01]])
    nine_poses = np.array([[-411.42, -1028.35],
                      [-460.69, -1062.27],
                      [-510.16, -1096.22],
                      [-468.21, -945.81],
                      [-517.41, -979.89],
                      [-566.79, -1013.61],
                      [-524.83, -863.37],
                      [-573.92, -897.33],
                      [-623.53, -931.27]])
    
    poses_gripper = nine_poses - base_pose

    A = np.ones(shape=(3,9))
    B = np.ones(shape=(3,9))

    A[:2, :] = pixels.T
    B[:2, :] = poses_gripper.T

    # 情况1: 当A @ A_t可逆时
    print(cali.solve_inv(A, B))

    # 情况2: 当A @ A_t不可逆时
    print(cali.solve_lstsq(A, B))

def test_CalibrateNinePointsTakeNinePhotos():
    cali = CalibrateNinePointsTakeNinePhotos()

    data_path = "D:/VSCodeProjects/calibration/data/eyeinhand_nine_points"
    pixels = cali.load_mark_pixels(data_path)

    base_pose = np.array([[-310.04, -659.13]])
    nine_poses = np.array([[-310.04, -659.13],
                            [-259.52, -709.13],
                            [-209.18, -783.75],
                            [-359.66, -678.15],
                            [-299.54, -766.49],
                            [-252.71, -824.54],
                            [-400.78, -711.28],
                            [-358.75, -763.67],
                            [-318.76, -858.00]])
    poses_gripper = nine_poses - base_pose

    A = np.ones(shape=(3,9))
    B = np.ones(shape=(3,9))

    A[:2, :] = pixels.T
    B[:2, :] = poses_gripper.T

    # 情况1: 当A @ A_t可逆时
    print(cali.solve_inv(A, B))

    # 情况2: 当A @ A_t不可逆时
    print(cali.solve_lstsq(A, B))

    # 情况1: 当A @ A_t可逆时
    # A_t = A.T
    # M = B @ A_t @ np.linalg.inv(A @ A_t)
    # print(M)

    # # 情况2: 当A @ A_t不可逆时
    # a_list = []
    # b_list = []
    # for i in range(9):
    #     a, b = special_resize(A[:,i], B[:,i])
    #     a_list.append(a)
    #     b_list.append(b)
    # A_ = np.concatenate(a_list, axis=0)
    # B_ = np.concatenate(b_list, axis=0)

    # print(np.linalg.lstsq(A_, B_, rcond=None)[0])



if __name__ == "__main__":
    # test_CalibrateNinePointsTakeAPhoto()
    test_CalibrateNinePointsTakeNinePhotos()
    
